import os
import pickle
from easydict import EasyDict
import logging
logger = logging.getLogger(__name__)
from utils.dirs import create_dirs


def save_cached_data(config, data_to_save, data_name, data_path=''):

    # Create the folder if not exist
    if not os.path.exists(config.cache.default_folder):
        create_dirs([config.cache.default_folder])

    if not data_path:
        data_path = os.path.join(
            config.cache.default_folder,
            '{}.pkl'.format(data_name))
    else:
        data_path = os.path.join(data_path, '{}.pkl'.format(data_name))

    with open(data_path, "wb" ) as f:
        logger.info('saving preprocessed data...')
        dump_data = {
            'cache': data_to_save,
        }
        pickle.dump(dump_data, f)
        logger.info(f'preprocessed data has been saved to {data_path}')



def load_cached_data(config, data_name, data_path='', condition=True):
    """[summary]

    Args:
        config ([type]): [description]
        data_name ([type]): [description]
        data_path (str, optional): [description]. Defaults to ''.
        condition (bool, optional): some condition to decide whether to load the cache. Defaults to True.

    Returns:
        [type]: [description]
    """
    if not data_path:
        data_path = os.path.join(
            config.cache.default_folder,
            '{}.pkl'.format(data_name))
    else:
        data_path = os.path.join(data_path, '{}.pkl'.format(data_name))
    
    # Check if config indicates to regenerate
    if config.cache.regenerate.get(data_name) is not None:
        if config.cache.regenerate[data_name]:
            logger.info('Data "{}" is forced to be re-generated by config file.'.format(data_name))
            return None
    else:
        logger.info('Data "{}" is not defined in config file.'.format(data_name))
        return None

    if os.path.exists(data_path) and condition:
        try:
            # Read data instead of re-generate
            logger.info(f'reading preprocessed data from {data_path}')
            with open(data_path, "rb" ) as f:
                load_pickle_data = pickle.load(f)['cache']
                return EasyDict(load_pickle_data)
        except Exception as e:
            logger.error('Failed to load pre-processed data, skipping...')
            logger.error(str(e))
    else:
        # This data is not cached
        return None